import argparse
import os

try:
    import hmark.parseutility as parser
except ImportError:
    import tools.parseutility as parser
import config
import logging
from tools.file_helper import get_all_file_paths, filepath_ends_in, get_files
import time
from concurrent.futures import ThreadPoolExecutor
from find_dep import FuncVulnContext, FuncVulnQueryContext


def generate_args():
    parser_ = argparse.ArgumentParser(
        description='A extractor used to extract all clang functions in specified file dir or .c/.h files.')
    parser_.add_argument('-f', '--input_files', dest='input_files', metavar='INPUT_FILE', nargs='*',
                         help='Specifies which files will be used to extract.')
    parser_.add_argument('-d', '--input_dir', dest='input_dir', metavar='INPUT_DIR', nargs='?',
                         help='Specifies which dir will be used to extract.')
    parser_.add_argument('-a', dest='abstraction_level', type=int, default=5, choices=range(0, 6),
                         help='Specify abstraction level, default lv5.')
    parser_.add_argument('-N', dest='use_normalization', action='store_true',
                         help='Normalize the function result extracted.')
    parser_.add_argument('--abs_base_dir', dest='abs_base_dir', metavar='ABS_BASE_DIR', default='',
                         help='Specifies which dir will be used to get changes and features for vuln functions.')

    def check(args):
        if not args.input_files and not args.input_dir:
            raise RuntimeError('Plz specify one of options -f or -d.')
        return args

    return check(parser_.parse_args())


if __name__ == '__main__':
    config.conf_log()
    args = generate_args()
    input_files = get_files(args.input_files, args.input_dir, config.VALID_SOURCE_SUFFICES)
    abstraction_level = args.abstraction_level
    use_normalization = args.use_normalization
    abs_base_dir = args.abs_base_dir

    if abs_base_dir:
        config.ABS_BASE_DIR = abs_base_dir
    abs_dir = os.path.join(config.ABS_BASE_DIR, 'abs')
    abs_files = parser.load_abs_feature_files(abs_dir)

    pool = ThreadPoolExecutor()
    context = FuncVulnQueryContext(FuncVulnContext(), pool, abstraction_level, use_normalization)

    try:
        logging.info('Starts to load all vuln funcs info into sys.')
        vuln_funcs_num = len(abs_files)
        context.load_vuln_func_features(abs_files, config.ENCODING)
        logging.info('Done, has loaded %s vuln funcs in total.' % vuln_funcs_num)

        start = time.time()
        test_func_num, test_lines, ret = context.query_files(input_files)
        end = time.time()

        logging.info('Total %s needed-test funcs have been processed.', test_func_num)
        logging.info('Total %s lines code has been processed.', test_lines)
        logging.info('Cost %ss in total.' % (end - start))
    finally:
        pool.shutdown()
    for _ in ret:
        print(_)
